package com.wondersgroup.aida.embeddingdemo;

import com.wondersgroup.aida.tools.ModelUtil;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.redis.RedisEmbeddingStore;

import java.util.List;

/**
 * @author: 紫麒麟vip
 * @create: 2025/2/16
 * Description:
 */
public class VectorDemo {
    public static void main(String[] args) {
        OpenAiEmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .apiKey(ModelUtil.API_KEY_0PENAI)
                .baseUrl(ModelUtil.BASE_URI_OPENAI)
                .build();

        //生成向量
        Response<Embedding> embed = embeddingModel.embed("你好，我是jszhao");
        System.out.println(embed.content().toString());
        System.out.println(embed.content().vector().length);

        // 我的名字叫jszhao
        RedisEmbeddingStore embeddingStore = RedisEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6379)
                // 维度，需要与计算结果保持一致。如果使用其他的模型，可能会有不同的结果
                .dimension(1536)
                .build();

        //存储向量--单独指定数据ID，索引的ID是固定的embedding-index
//        embeddingStore.add("vecl", embed.content());

        // 计算"我的名字叫jszhao"与向量数据库中存储的"你好，我是jszhao"的向量相似度
        List<EmbeddingMatch<TextSegment>> matches = embeddingStore.findRelevant(embeddingModel.embed("我的名字叫jszhao").content(), 3, -1);
        for (EmbeddingMatch<TextSegment> match : matches) {
            System.out.println(match.score());
        }
    }
}
